#!/usr/bin/env bash

. ./path.sh || exit 1;
. ./cmd.sh || exit 1;

#. utils/parse_options.sh || exit 1;

# general configuration
stage=1       # now can not set 0
stop_stage=6
cmd=run.pl

backend=pytorch
ngpu=4         # number of gpus ("0" uses cpu, otherwise use gpu)
debugmode=1
dumpdir=dump   # directory to dump full features
verbose=0      # verbose option

# feature configuration
do_delta=false
online_cmvn_opts=" --config=conf/online_cmvn.conf "
compress=true
global_cmvn_file=
add_deltas=false


# Set bash to 'debug' mode, it will exit on :
# -e 'error', -u 'undefined variable', -o ... 'error in pipeline', -x 'print commands',
set -e
set -u
set -o pipefail

if [ $# != 3 ]; then
  echo "Usage: "
  echo "  $0 <data-path> <ali-dir> <dst-dir>"
  echo "e.g.:"
  echo "  $0 data/traindata exp/mono0a exp/data2json"
  exit 1;
fi

data=$1
alidir=$2
dir=$3

for f in $data/feats.scp $alidir/ali.1.gz $alidir/final.mdl $alidir/tree; do
  [ ! -f $f ] && echo "$0: no such file $f" && exit 1;
done

nj=$(cat $alidir/num_jobs)
sdata=$data/split$nj
if [ ! -d $sdata ]; then
  utils/split_data.sh $data $nj
fi

mkdir -p $dir/log $dir/info
cp $alidir/tree $dir
cp $alidir/num_jobs $dir/info


if [ -e $data/cmvn.scp ]; then
  cmvn_scp=$data/cmvn.scp
  if [ -e $data/global_cmvn.stats ]; then
    global_cmvn_file=$data/global_cmvn.stats
  else
    matrix-sum --binary=false scp:$cmvn_scp -> $dir/global_cmvn.stats
	global_cmvn_file=$dir/global_cmvn.stats
  fi
else
  cmvn_scp="none"
  global_cmvn_file="none"
fi



if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
	applied_cmvn_data=$dir/online_cmvn_data
	[  -d $applied_cmvn_data ] && rm -rf $applied_cmvn_data
	utils/copy_data_dir.sh $data $applied_cmvn_data
	[ -f $applied_cmvn_data/feats.scp ] && rm -rf $applied_cmvn_data/feats.scp
	mkdir $applied_cmvn_data/log
	echo "$0: apply-online-cmvn"
	$cmd JOB=1:$nj $applied_cmvn_data/log/applied_cmvn_data.JOB.log \
	apply-cmvn-online $online_cmvn_opts --spk2utt=ark:$sdata/JOB/spk2utt $global_cmvn_file  scp:$sdata/JOB/feats.scp ark,scp:$sdata/JOB/applied_cmvn_data.JOB.ark,$sdata/JOB/applied_cmvn_data.JOB.scp
	
	if $do_delta;then
	$cmd JOB=1:$nj $applied_cmvn_data/log/add_deltas_data.JOB.log \
		add-deltas ark:$sdata/JOB/applied_cmvn_data.JOB.ark ark,scp:$sdata/JOB/add_deltas.JOB.ark,$sdata/JOB/add_deltas.JOB.scp
		for n in $(seq $nj); do
			rm -f $sdata/$n/applied_cmvn_data.$n.*
		done
		for n in $(seq $nj);do
			cat $sdata/$n/add_deltas.$n.scp || exit 1
		done >> $applied_cmvn_data/feats.scp
	else
		for n in $(seq $nj);do
			cat $sdata/$n/applied_cmvn_data.$n.scp || exit 1
		done >> $applied_cmvn_data/feats.scp
	fi
	
	
	utils/fix_data_dir.sh $applied_cmvn_data
	
	sdata=$applied_cmvn_data/split$nj
	if [ ! -d $sdata ]; then
	  utils/split_data.sh $applied_cmvn_data $nj
	fi
	
  echo "$0: working out feature length"
  $cmd JOB=1:$nj $dir/log/get_feat_len.JOB.log \
    feat-to-len scp:$sdata/JOB/feats.scp ark,t:$dir/log/feats_lengths.JOB || exit 1;

  for n in $(seq $nj); do
    cat $dir/log/feats_lengths.$n || exit 1
  done > $dir/feats_lengths || exit 1
  rm $dir/log/feats_lengths.*
fi



if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
  echo "$0: working out feature dim and num_pdfs"

  if feat_dim=$(feat-to-dim "scp:$data/feats.scp" - 2>/dev/null); then
    echo $feat_dim > $dir/info/feat_dim
  else # run without redirection to show the error.
    feat-to-dim "scp:$data/feats.scp" -; exit 1
  fi

  num_pdfs=$(tree-info --print-args=false $alidir/tree | grep num-pdfs | awk '{print $2}')
  echo $num_pdfs >$dir/info/num_pdfs
else
  feat_dim=$(cat $dir/info/feat_dim)
  num_pdfs=$(cat $dir/info/num_pdfs)
fi

if [ $stage -le 3 ] && [ ${stop_stage} -ge 3 ]; then
  echo "$0: making align pdf label scp and working out label length ."
  
  $cmd JOB=1:$nj $dir/log/mk_align_label.JOB.log \
    ali-to-pdf $alidir/final.mdl "ark:gunzip -c $alidir/ali.JOB.gz|" ark,t: \| \
    copy-int-vector ark,t:- ark,scp:$dir/ali_pdf.JOB.ark,$dir/ali_pdf.JOB.scp || exit 1;
	
	$cmd JOB=1:$nj $dir/log/mk_align_label.JOB.log \
	ali-to-pdf $alidir/final.mdl "ark:gunzip -c $alidir/ali.JOB.gz|" ark:- \| \
	ali-to-post ark:- ark,t:$dir/ali_post.JOB.ark.txt || exit 1;
	

  for n in $(seq $nj); do
    cat $dir/ali_pdf.$n.scp || exit 1
    #rm $dir/ali_pdf.$n.scp
  done > $dir/ali_pdf.scp || exit 1
  
  
  # for n in $(seq $nj); do
    # cat $dir/ali_post.$n.scp || exit 1
    # #rm $dir/ali_pdf.$n.scp
  # done > $dir/ali_post.scp || exit 1
  
  
  for n in $(seq $nj); do
    copy-int-vector ark:$dir/ali_pdf.$n.ark ark,t: | awk '{print $1" "NF-1}' || exit 1
  done > $dir/ali_pdf_lengths || exit 1

fi

echo " all needed data group into a json in " $applied_cmvn_data

#python utils/align_label_data2json.py $applied_cmvn_data/feats.scp $dir/feats_lengths $feat_dim $dir/ali_pdf.scp $dir/ali_pdf_lengths $num_pdfs $applied_cmvn_data/utt2spk $cmvn_scp $global_cmvn_file > $dir/data.json
#python utils/align_label_data2json.py $applied_cmvn_data/feats.scp $dir/feats_lengths $feat_dim $dir/ali_pdf.scp $dir/ali_pdf_lengths $num_pdfs $data/utt2spk $cmvn_scp $global_cmvn_file > $dir/data.json
